#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient
from bs4 import BeautifulSoup
import os
import multiprocessing
import sys
from tqdm import tqdm


# 连接MongoDB
def InitMongoDB():
    MONGODB_HOST = 'localhost'
    MONGODB_PORT = 27017
    # MONGODB_DB_NAME = 'sra_stat'
    # MONGODB_COLLECTION_NAME = 'sra_stat'
    connection = MongoClient(MONGODB_HOST, MONGODB_PORT)
    # db = connection[MONGODB_DB_NAME]
    # collection = db[MONGODB_COLLECTION_NAME]
    return connection


# 关闭MongoDB
def CloseMongoDB(MongoDB_Connection):
    MongoDB_Connection.close()


def Stat_File(file_path, dirname, collection):
    filename = file_path.split("/")[-1]
    if filename.find("submission") != -1:
        collection.update_one({"id": dirname}, {"$inc": {"submission_file": 1}})
        submission_count = Count_Submission(file_path)
        if submission_count > 0:
            collection.update_one({"id": dirname}, {"$inc": {"submission": submission_count}})
    elif filename.find("study") != -1:
        collection.update_one({"id": dirname}, {"$inc": {"study_file": 1}})
        study_count = Count_Study(file_path)
        if study_count > 0:
            collection.update_one({"id": dirname}, {"$inc": {"study": study_count}})
    elif filename.find("experiment") != -1:
        collection.update_one({"id": dirname}, {"$inc": {"experiment_file": 1}})
        experiment_count = Count_Experiment(file_path)
        if experiment_count > 0:
            collection.update_one({"id": dirname}, {"$inc": {"experiment": experiment_count}})
    elif filename.find("sample") != -1:
        collection.update_one({"id": dirname}, {"$inc": {"sample_file": 1}})
        sample_count = Count_Sample(file_path)
        if sample_count > 0:
            collection.update_one({"id": dirname}, {"$inc": {"sample": sample_count}})
    elif filename.find("run") != -1:
        collection.update_one({"id": dirname}, {"$inc": {"run_file": 1}})
        run_count = Count_Run(file_path)
        if run_count > 0:
            collection.update_one({"id": dirname}, {"$inc": {"run": run_count}})
    elif filename.find("analysis") != -1:
        collection.update_one({"id": dirname}, {"$inc": {"analysis_file": 1}})
        analysis_count = Count_Analysis(file_path)
        if analysis_count > 0:
            collection.update_one({"id": dirname}, {"$inc": {"analysis": analysis_count}})


def Count_Submission(file_path):
    soup = BeautifulSoup(open(file_path), 'lxml-xml')
    return len(soup.find_all("SUBMISSION"))


def Count_Study(file_path):
    soup = BeautifulSoup(open(file_path), 'lxml-xml')
    return len(soup.find_all("STUDY"))


def Count_Experiment(file_path):
    soup = BeautifulSoup(open(file_path), 'lxml-xml')
    return len(soup.find_all("EXPERIMENT"))


def Count_Sample(file_path):
    soup = BeautifulSoup(open(file_path), 'lxml-xml')
    return len(soup.find_all("SAMPLE"))


def Count_Run(file_path):
    soup = BeautifulSoup(open(file_path), 'lxml-xml')
    return len(soup.find_all("RUN"))


def Count_Analysis(file_path):
    soup = BeautifulSoup(open(file_path), 'lxml-xml')
    return len(soup.find_all("ANALYSIS"))


def Stat_Dir(dir, dirname):
    client = InitMongoDB()
    collection = client["sra_stat"]["sra_stat"]
    for filename in os.listdir(dir):
        file_path = os.path.join(dir, filename)
        Stat_File(file_path, dirname, collection)
    CloseMongoDB(client)


if __name__ == "__main__":
    dir = sys.argv[1]
    parent_dir = os.path.abspath(dir)
    dirname = os.path.basename(parent_dir)
    print "Start handling " + dirname + " ..."
    client = InitMongoDB()
    collection = client["sra_stat"]["sra_stat"]
    collection.update_one({"id": dirname}, {
        "$set": {"submission_file": 0, "submission": 0, "study_file": 0, "study": 0, "experiment_file": 0,
                 "experiment": 0, "sample_file": 0, "sample": 0, "run_file": 0, "run": 0, "analysis_file": 0,
                 "analysis": 0}}, upsert=True)
    CloseMongoDB(client)
    pool = multiprocessing.Pool(multiprocessing.cpu_count())
    list = os.listdir(parent_dir)
    for item in tqdm(list):
        dirpath = os.path.join(parent_dir, item)
        if os.path.isdir(dirpath):
            pool.apply_async(Stat_Dir, (dirpath, dirname,))
    pool.close()
    pool.join()
    print "complete."
